package org.apache.nifi.processors.huawei.smn;

import com.huaweicloud.sdk.smn.v2.SmnClient;
import com.huaweicloud.sdk.smn.v2.model.PublishMessageRequest;
import com.huaweicloud.sdk.smn.v2.model.PublishMessageResponse;
import org.apache.nifi.processors.huawei.abstractprocessor.AbstractHuaweiProcessor;
import org.apache.nifi.processors.huawei.common.PropertyDescriptors;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;

import java.util.HashMap;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;

public class TestPublishSMNMessage {

    private TestRunner runner = null;
    private PublishSMNMessage mockPublishSMNMessage = null;
    private SmnClient mockSmnClient = null;

    @BeforeEach
    public void setUp() {

        mockSmnClient = Mockito.mock(SmnClient.class);
        mockPublishSMNMessage = new PublishSMNMessage() {
            @Override
            public SmnClient getSmnClient() { return mockSmnClient; }
        };
        runner = TestRunners.newTestRunner(mockPublishSMNMessage);
    }

    @Test
    public void testPublishMessage() {

        runner.setProperty(PropertyDescriptors.ACCESS_KEY, "test-ak");
        runner.setProperty(PropertyDescriptors.SECRET_KEY, "test-sk");
        runner.setProperty(PropertyDescriptors.SMN_REGION, "ap-southeast-3");
        runner.setProperty(PublishSMNMessage.SMN_TOPIC_URN, "urn:smn:ap-southeast-3:test-projectId:AirFlowShowCase");
        runner.setProperty(PublishSMNMessage.SMN_PROJECT_ID, "test-projectId");
        runner.setProperty(PublishSMNMessage.SMN_SUBJECT, "${eval.subject}");
        runner.setProperty(PublishSMNMessage.SMN_MESSAGE, "${eval.message}");
        runner.setValidateExpressionUsage(false);

        final Map<String, String> flowFileAttributes = new HashMap<>();
        flowFileAttributes.put("eval.subject", "test-subject");
        flowFileAttributes.put("eval.message", "test-message");
        runner.enqueue("test-message", flowFileAttributes);

        PublishMessageResponse mockPublishMessageResponse = new PublishMessageResponse();
        Mockito.when(mockSmnClient.publishMessage(Mockito.any(PublishMessageRequest.class))).thenReturn(mockPublishMessageResponse);

        runner.run();

        ArgumentCaptor<PublishMessageRequest> captureRequest = ArgumentCaptor.forClass(PublishMessageRequest.class);
        Mockito.verify(mockSmnClient, Mockito.times(1)).publishMessage(captureRequest.capture());
        PublishMessageRequest request = captureRequest.getValue();
        assertEquals("urn:smn:ap-southeast-3:test-projectId:AirFlowShowCase", request.getTopicUrn());
        assertEquals("test-message", request.getBody().getMessage());
        assertEquals("test-subject", request.getBody().getSubject());

        runner.assertAllFlowFilesTransferred(AbstractHuaweiProcessor.REL_SUCCESS, 1);
    }

    @Test
    public void testPublishMessageFailure() {

        runner.setProperty(PropertyDescriptors.ACCESS_KEY, "test-ak");
        runner.setProperty(PropertyDescriptors.SECRET_KEY, "test-sk");
        runner.setProperty(PropertyDescriptors.SMN_REGION, "ap-southeast-3");
        runner.setProperty(PublishSMNMessage.SMN_TOPIC_URN, "urn:smn:ap-southeast-3:test-projectId:AirFlowShowCase");
        runner.setProperty(PublishSMNMessage.SMN_PROJECT_ID, "test-projectId");
        runner.setProperty(PublishSMNMessage.SMN_SUBJECT, "${eval.subject}");
        runner.setProperty(PublishSMNMessage.SMN_MESSAGE, "${eval.message}");
        runner.setValidateExpressionUsage(false);

        final Map<String, String> flowFileAttributes = new HashMap<>();
        flowFileAttributes.put("eval.subject", "test-subject");
        flowFileAttributes.put("eval.message", "test-message");
        runner.enqueue("test-message", flowFileAttributes);

        Mockito.when(mockSmnClient.publishMessage(Mockito.any(PublishMessageRequest.class))).thenThrow(new RuntimeException());

        runner.run();

        ArgumentCaptor<PublishMessageRequest> captureRequest = ArgumentCaptor.forClass(PublishMessageRequest.class);
        Mockito.verify(mockSmnClient, Mockito.times(1)).publishMessage(captureRequest.capture());

        runner.assertAllFlowFilesTransferred(AbstractHuaweiProcessor.REL_FAILURE, 1);
    }
}
